#ifndef UHM_SOLVER_OPS_MERGE_HXX_
#define UHM_SOLVER_OPS_MERGE_HXX_

namespace uhm {

  template<class S_> Int_ op_uhm_merge_A(S_ *s, typename S_::_Tree_::Tnd_* me) {

    typename S_::_Tree_::_Data_ elt = me->get_data();
    typename S_::_Mesh_::_Uhm_ &uhm = elt->uhm();
    
    // As default, ABR is not preallocated
    if (uhm.ABR().is_flat_buffer_null()) 
      uhm.ABR().create_flat_buffer();
    
    if (me->is_leaf()) { 
      // update from work cache
      Map_< typename S_::_Mesh_::_Id_Nd_> factor, schur;
      elt->nds2map(factor, schur);

      typename S_::_Cache_::Entry_ in(factor, schur), *out;
      UHM_ERROR(s->work().pull(in, &out), UHM_ERROR_SOLVER_PULLED_WORK);

      Mapper_< typename S_::_Mesh_::_Id_Nd_ >
        m_s2f(&out->schur(), &elt->factor()),
        m_s2s(&out->schur(), &elt->schur() );

      m_s2f.relate();
      m_s2s.relate();

      // Mapping A
      switch(s->get_uplo()) {
      case Mat_Base_::FULL:
        copy< typename S_::_Mesh_::_Id_Nd_ >(m_s2f, m_s2s, out->uhm().ABR(), uhm.ATR());
      case Mat_Base_::LOWER_TRIANGULAR:
        copy< typename S_::_Mesh_::_Id_Nd_ >(m_s2f, m_s2f, out->uhm().ABR(), uhm.ATL());
        copy< typename S_::_Mesh_::_Id_Nd_ >(m_s2s, m_s2f, out->uhm().ABR(), uhm.ABL());
        copy< typename S_::_Mesh_::_Id_Nd_ >(m_s2s, m_s2s, out->uhm().ABR(), uhm.ABR());
      }

    } else {

      // update from children
      for (Int_ i=0;i<me->get_n_children();++i) {
        
        typename S_::_Tree_::_Data_ c_elt = me->get_child(i)->get_data();
        typename S_::_Mesh_::_Uhm_ &c_uhm = c_elt->uhm();
        
        Mapper_< typename S_::_Mesh_::_Id_Nd_ > m_s2f, m_s2s;
        c_elt->relation().get(m_s2f, m_s2s);

        Map_< ID_Node_H_<1> > c_s2f, e_s2f, c_s2s, e_s2s;
        m_s2f.coarse(c_s2f, e_s2f);
        m_s2s.coarse(c_s2s, e_s2s);

        Mapper_< ID_Node_H_<1> > s2f(&c_s2f, &e_s2f), s2s(&c_s2s, &e_s2s);
        s2f.relate();
        s2s.relate();

        // Mapping A
        switch(s->get_uplo()) {
        case Mat_Base_::FULL:
          axpy< ID_Node_H_<1> >(s2f, s2s, c_uhm.ABR(), uhm.ATR());
        case Mat_Base_::LOWER_TRIANGULAR:
          axpy< ID_Node_H_<1> >(s2f, s2f, c_uhm.ABR(), uhm.ATL());
          axpy< ID_Node_H_<1> >(s2s, s2f, c_uhm.ABR(), uhm.ABL());
          axpy< ID_Node_H_<1> >(s2s, s2s, c_uhm.ABR(), uhm.ABR());
        }

        // One A is formed, children schur's are not necessary.
        if (!c_uhm.ABR().is_flat_buffer_null()) 
          c_uhm.ABR().free_flat_buffer();;
      }
    }

    return true;
  }

  template<class S_> static Int_ op_uhm_merge_BXR(S_ *s, typename S_::_Tree_::Tnd_* me,
                                                  Int_ b, Int_ x, Int_ r);

  template<class S_> static Int_ op_uhm_merge_BXR(S_ *s, typename S_::_Tree_::Tnd_* me,
                                                  Int_ b, Int_ x, Int_ r) {

    // b   - used for solve_0 phase, setting a problem
    // x,r - trsm

    typename S_::_Tree_::_Data_ elt = me->get_data();
    typename S_::_Mesh_::_Uhm_ &uhm = elt->uhm();

    if (me->is_leaf()) {
      if (x || r==1) 
        return true;

      // Update to work cache
      Map_< typename S_::_Mesh_::_Id_Nd_> factor, schur;
      elt->nds2map(factor, schur);

      typename S_::_Cache_::Entry_ in(factor, schur), *out;
      UHM_ERROR(s->work().pull(in, &out), UHM_ERROR_SOLVER_PULLED_WORK);

      Mapper_< typename S_::_Mesh_::_Id_Nd_ >
        m_s2f(&out->schur(), &elt->factor()),
        m_s2s(&out->schur(), &elt->schur() );

      m_s2f.relate();
      m_s2s.relate();

      // Mapping X from B in work
      if (b) {
        copy< typename S_::_Mesh_::_Id_Nd_ >(m_s2f, out->uhm().BB(), uhm.BT());
        copy< typename S_::_Mesh_::_Id_Nd_ >(m_s2s, out->uhm().BB(), uhm.BB());

        copy(uhm.BT(), uhm.XT());
        copy(uhm.BB(), uhm.XB());
      }

      if (r==2) {
        copy< typename S_::_Mesh_::_Id_Nd_ >(m_s2f, out->uhm().RB(), uhm.RT());
        copy< typename S_::_Mesh_::_Id_Nd_ >(m_s2s, out->uhm().RB(), uhm.RB());
      }

    } else {

      if (b) {
        if (elt->factor().get_size()) {
          uhm.BT().set_zero();
          uhm.XT().set_zero();
        }
        if (elt->schur().get_size()){ 
          uhm.BB().set_zero();
          uhm.XB().set_zero();
        }
      }

      if (r) {
        if (elt->factor().get_size())
          uhm.RT().set_zero();
        if (elt->schur().get_size())
          uhm.RB().set_zero();
      }

      for (Int_ i=0;i<me->get_n_children();++i) {
        
        typename S_::_Tree_::_Data_ c_elt = me->get_child(i)->get_data();

        Mapper_< typename S_::_Mesh_::_Id_Nd_ > m_s2f, m_s2s;
        c_elt->relation().get(m_s2f, m_s2s);
        
        typename S_::_Mesh_::_Uhm_ &c_uhm = c_elt->uhm();


        Map_< ID_Node_H_<1> > c_s2f, e_s2f, c_s2s, e_s2s;
        m_s2f.coarse(c_s2f, e_s2f);
        m_s2s.coarse(c_s2s, e_s2s);

        Mapper_< ID_Node_H_<1> > s2f(&c_s2f, &e_s2f), s2s(&c_s2s, &e_s2s);
        s2f.relate();
        s2s.relate();

        // Mapping B 
        if (b) {
          axpy< ID_Node_H_<1> >(s2f, c_uhm.BB(), uhm.BT());
          axpy< ID_Node_H_<1> >(s2s, c_uhm.BB(), uhm.BB());
          c_uhm.BB().set_zero(); 
          c_uhm.XB().set_zero(); 
        }
        
        // Mapping X
        if (x) {
          axpy< ID_Node_H_<1> >(s2f, c_uhm.XB(), uhm.XT());
          axpy< ID_Node_H_<1> >(s2s, c_uhm.XB(), uhm.XB());
        }
        
        // Mapping R
        switch (r) {
        case 1:
          axpy< ID_Node_H_<1> >(s2f, c_uhm.RB(), uhm.RT());
          axpy< ID_Node_H_<1> >(s2s, c_uhm.RB(), uhm.RB());
          break;
        case 2:
          copy< ID_Node_H_<1> >(s2f, c_uhm.RB(), uhm.RT());
          copy< ID_Node_H_<1> >(s2s, c_uhm.RB(), uhm.RB());
          c_uhm.RB().set_zero();
          break;
        }
      }

      if (b) {
        copy(uhm.BT(), uhm.XT());
        copy(uhm.BB(), uhm.XB());
      }
    }
    return true;
  }

  template<class S_> Int_ op_uhm_merge_B(S_ *s, typename S_::_Tree_::Tnd_* me) {
    return op_uhm_merge_BXR(s, me, 1, 0, 0);
  }
  template<class S_> Int_ op_uhm_merge_X(S_ *s, typename S_::_Tree_::Tnd_* me) {
    return op_uhm_merge_BXR(s, me, 0, 1, 0);
  }
  template<class S_> Int_ op_uhm_merge_R(S_ *s, typename S_::_Tree_::Tnd_* me) {
    return op_uhm_merge_BXR(s, me, 0, 0, 1);
  }
  template<class S_> Int_ op_uhm_copy_R(S_ *s, typename S_::_Tree_::Tnd_* me) {
    return op_uhm_merge_BXR(s, me, 0, 0, 2);
  }

}


#endif

